# -*- coding : utf-8 -*-


import time
import json
import torch
import logging
import transformers
import tornado.httpclient

tokenizer = transformers.AutoTokenizer.from_pretrained("../bloom-1.7b-credit-v2", padding_side="right")
model = transformers.AutoModelForCausalLM.from_pretrained("../bloom-1.7b-credit-v2")

model = model.to("cuda:1")

class dialogueDataHandler(tornado.web.RequestHandler):

    def generate_prompt(self, instruction, input_data=None):
        if input_data:
            return f"""Instruction:\n{instruction}\n\n### Input:{input_data}\n\n### Response:"""
        else:
            return f"""Instruction:\n{instruction}\n\n### Response:"""


    def run(self, instruction, input_data):

        prompt = self.generate_prompt(instruction, input_data)
        inputs = tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to("cuda:1")

        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                temperature=0.8,
                top_p=0.95,
                do_sample=True,
                num_beams=1,
                max_new_tokens=1000,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.pad_token_id,
                return_dict_in_generate=True,
                output_scores=True
            )

            s = generation_output.sequences[0]
            output = tokenizer.decode(s)
            res = output.split("### Response:")[-1].strip().replace('</s>', '')
            logging.info('输出：{}'.format(res))
            return res


    def get(self):
        instruction = self.get_argument('instruction', '')
        input_data = self.get_argument('input', '')
        logging.info('instruction:{}\ninput:{}'.format(instruction, input_data))

        output_data = self.run(instruction, input_data)

        self.write(output_data)


    def post(self):

        instruction = self.get_argument('instruction', '')
        input_data = self.get_argument('input', '')
        logging.info('instruction:{}\ninput:{}'.format(instruction, input_data))

        output_data = self.run(instruction, input_data)

        self.write(output_data)




if __name__ == '__main__':


    tornado.options.parse_command_line()
    app = tornado.web.Application(handlers=[
        (r"/credit_v2", dialogueDataHandler),
    ], autoreload=False, debug=False)
    http_server = tornado.httpserver.HTTPServer(app)
    http_server.bind(8200)
    http_server.start(2)
    tornado.ioloop.IOLoop.current().start()

    logging.info("accuracy server start")






